import numpy as np
import matplotlib.pyplot as plt

# 显示图像的函数


def show_image():
    plt.figure(figsize=(8, 6))
    plt.scatter(x, y_noisy, label='Noisy Data')
    plt.plot(x, a_fit * x + b_fit, color='red', label='Fitted Line')
    plt.plot(x, y_true, color='green', linestyle='--', label='True Line')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('Linear Regression with Noisy Data')
    plt.legend()
    plt.grid(True)
    plt.show()

# 生成带噪声的样本数据


def generate_data_set():
    np.random.seed(0)  # 设置随机种子以确保结果可重现
    a_true = 2.1       # 真实参数a
    b_true = 3         # 真实参数b
    num_samples = 50   # 样本数量
    noise = np.random.normal(0, 1, num_samples)  # 生成（0， 1）正态分布的随机数充当噪声
    x = np.random.uniform(0, 10, num_samples)    # 生成分布均匀的随机数作为x示例
    y_true = a_true * x + b_true                 # 获取真实数据的标签
    y_noisy = y_true + noise                     # 获取添加噪声之后的样本数据集
    return (x, y_noisy, y_true)

# 最小二乘法拟合


def least_squares_fit():
    X = np.column_stack((x, np.ones_like(x)))       # 构建设计矩阵，添加常数列，将a和b吸收成向量形式
    theta = np.linalg.inv(X.T @ X) @ X.T @ y_noisy  # 最小二乘法求解参数
    a_fit, b_fit = theta                            # 提取拟合得到的斜率和截距
    return (a_fit, b_fit)


if __name__ == "__main__":
    # 生成数据集
    x, y_noisy, y_true = generate_data_set()
    # 使用最小二乘法求解线性回归参数
    a_fit, b_fit = least_squares_fit()
    # 绘制样本点和拟合直线
    show_image()
